# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np

# sklearn / onnx
from sklearn.ensemble import IsolationForest
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

# SQLAlchemy（推荐）
from sqlalchemy import create_engine, text

# ===== 1) 连接数据库（用 SQLAlchemy） =====
# 安装依赖：pip install mysql-connector-python SQLAlchemy
DB_USER = "root"
DB_PWD  = "123456"
DB_HOST = "192.168.140.128"
DB_PORT = 3306
DB_NAME = "yuapi"

# 使用 mysql-connector-python 的 driver
engine = create_engine(
    f"mysql+mysqlconnector://{DB_USER}:{DB_PWD}@{DB_HOST}:{DB_PORT}/{DB_NAME}?charset=utf8mb4",
    pool_pre_ping=True
)

# ===== 2) 读取 gateway_access_log（最近 30 天） =====
query = text("""
SELECT header_cnt, param_cnt, query_size, body_size,
       method, ak, user_id, success, latency_ms
FROM gateway_access_log
WHERE created_at >= NOW() - INTERVAL 30 DAY
""")
df = pd.read_sql(query, engine)

if df.empty:
    raise RuntimeError("gateway_access_log 没有数据，先导入或生成一些再训练。")

# ===== 3) 特征工程（与 Java ORDER 对齐）=====
df["method_get"] = (df["method"].str.upper() == "GET").astype(int)
df["method_post"] = (df["method"].str.upper() == "POST").astype(int)
df["has_ak"]   = df["ak"].notna().astype(int)
df["has_user"] = df["user_id"].notna().astype(int)

FEATURES = [
    "header_cnt","param_cnt","query_size","body_size",
    "method_get","method_post",
    "has_ak","has_user",
    "success","latency_ms"
]

X = df[FEATURES].fillna(0).astype(float).values

# ===== 4) 训练 IsolationForest（无监督）=====
pipe = Pipeline([
    ("scaler", StandardScaler()),
    ("iforest", IsolationForest(
        n_estimators=150,
        contamination=0.05,   # 估计异常占比，可根据数据调
        random_state=42
    ))
])
pipe.fit(X)

# ===== 5) 导出 ONNX（关键：指定 ai.onnx.ml 的 opset=3）=====
initial_type = [("input", FloatTensorType([None, len(FEATURES)]))]
onnx_model = convert_sklearn(
    pipe,
    initial_types=initial_type,
    target_opset={ "": 13, "ai.onnx.ml": 3 }   # 修复你的报错
)

with open("iforest.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

print("✅ 模型导出成功: iforest.onnx")
print("特征顺序（务必与 Java ORDER 一致）:", FEATURES)
